{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "3YLyiTwPfVCT"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/magenta/ddsp/blob/main/ddsp/colab/demos/pitch_detection.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/\u003e\u003c/a\u003e\n",
        "\n",
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Bvp6GWqtfVCW"
      },
      "outputs": [],
      "source": [
        "# Copyright 2020 Google LLC. All Rights Reserved.\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JndnmDMp66FL"
      },
      "source": [
        "# DDSP Pitch Detection Demo\n",
        "\n",
        "This notebook is a demo of pitch detection using inverse audio synthesis. \n",
        "\n",
        "\n",
        "* [ICML Workshop paper](https://openreview.net/forum?id=RlVTYWhsky7)\n",
        "* [Audio Examples](http://goo.gl/magenta/ddsp-inv) \n",
        "\n",
        "This notebook extracts these features from input audio (either uploaded files, or recorded from the microphone) and resynthesizes the audio from the model. The DDSP-INV model is hierarchical, and provides both resynthesis from the sinusoidal model and harmonic model.\n",
        "\n",
        "\u003cimg src=\"https://storage.googleapis.com/ddsp-inv/full_stack/diagram.png\" alt=\"DDSP Pitch Detection\" width=\"700\"\u003e\n",
        "\n",
        "\n",
        "### Instructions for running:\n",
        "\n",
        "* Make sure to use a GPU runtime, click:  __Runtime \u003e\u003e Change Runtime Type \u003e\u003e GPU__\n",
        "* Press ▶️ on the left of each of the cells\n",
        "* View the code: Double-click any of the cells\n",
        "* Hide the code: Double click the right side of the cell\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "6wZde6CBya9k"
      },
      "outputs": [],
      "source": [
        "#@title #Install and Import\n",
        "\n",
        "#@markdown Install ddsp, define some helper functions, and download the model. This transfers a lot of data and _should take a minute or two_.\n",
        "%tensorflow_version 2.x\n",
        "print('Installing from pip package...')\n",
        "!pip install -qU ddsp==0.13.0\n",
        "\n",
        "# Ignore a bunch of deprecation warnings\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "import copy\n",
        "import os\n",
        "import time\n",
        "\n",
        "import ddsp\n",
        "import ddsp.training\n",
        "from ddsp.colab import colab_utils\n",
        "from ddsp.colab.colab_utils import (play, record, \n",
        "    specplot, upload, DEFAULT_SAMPLE_RATE)\n",
        "import gin\n",
        "from google.colab import files\n",
        "import librosa\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "# Helper Functions\n",
        "sample_rate = DEFAULT_SAMPLE_RATE  # 16000\n",
        "\n",
        "\n",
        "print('Done!')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "Go36QW9AS_CD"
      },
      "outputs": [],
      "source": [
        "#@title Record or Upload Audio\n",
        "#@markdown * Either record audio from microphone or upload audio from file (.mp3 or .wav) \n",
        "#@markdown * Audio should be monophonic (single instrument / voice)\n",
        "\n",
        "record_or_upload = \"Upload (.mp3 or .wav)\"  #@param [\"Record\", \"Upload (.mp3 or .wav)\"]\n",
        "\n",
        "record_seconds =     5#@param {type:\"number\", min:1, max:10, step:1}\n",
        "\n",
        "if record_or_upload == \"Record\":\n",
        "  audio = record(seconds=record_seconds)\n",
        "else:\n",
        "  # Load audio sample here (.mp3 or .wav3 file)\n",
        "  # Just use the first file.\n",
        "  filenames, audios = upload()\n",
        "  audio = audios[0]\n",
        "audio = audio[np.newaxis, :]\n",
        "\n",
        "\n",
        "# Plot.\n",
        "specplot(audio)\n",
        "play(audio)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "wmSGDWM5yyjm"
      },
      "outputs": [],
      "source": [
        "#@title Load a model\n",
        "#@markdown Run for every new audio input. Models separately trained on the [URMP](http://www2.ece.rochester.edu/projects/air/projects/URMP/annotations_5P.html), [MDB-stem-synth](https://zenodo.org/record/1481172#.Xzouy5NKhTY), and [MIR1k](https://sites.google.com/site/unvoicedsoundseparation/mir-1k) datasets.\n",
        "model = 'urmp' #@param ['urmp', 'mdb_stem_synth', 'mir1k']\n",
        "MODEL = model\n",
        "\n",
        "# Pretrained models.\n",
        "PRETRAINED_DIR = '/content/pretrained'\n",
        "# Copy over from gs:// for faster loading.\n",
        "!rm -r $PRETRAINED_DIR \u0026\u003e /dev/null\n",
        "!mkdir $PRETRAINED_DIR \u0026\u003e /dev/null\n",
        "GCS_CKPT_DIR = 'gs://ddsp-inv/ckpts'\n",
        "model_dir = os.path.join(GCS_CKPT_DIR, '%s_ckpt' % model.lower())\n",
        "\n",
        "!gsutil cp $model_dir/* $PRETRAINED_DIR \u0026\u003e /dev/null\n",
        "model_dir = PRETRAINED_DIR\n",
        "gin_file_pattern = os.path.join(model_dir, 'operative_config*.gin')\n",
        "gin_file = tf.io.gfile.glob(gin_file_pattern)[0]\n",
        "\n",
        "\n",
        "# Parse gin config,\n",
        "with gin.unlock_config():\n",
        "  gin.parse_config_file(gin_file, skip_unknown=True)\n",
        "\n",
        "# Assumes only one checkpoint in the folder, 'ckpt-[iter]`.\n",
        "ckpt_files = [f for f in tf.io.gfile.listdir(model_dir) if 'ckpt' in f]\n",
        "ckpt_name = ckpt_files[0].split('.')[0]\n",
        "ckpt = os.path.join(model_dir, ckpt_name)\n",
        "\n",
        "# Ensure dimensions and sampling rates are equal\n",
        "time_steps_train = 125\n",
        "n_samples_train = 64000\n",
        "hop_size = int(n_samples_train / time_steps_train)\n",
        "\n",
        "time_steps = int(audio.shape[1] / hop_size)\n",
        "n_samples = time_steps * hop_size\n",
        "audio = audio[:, :n_samples]\n",
        "\n",
        "gin_params = [\n",
        "    'TranscribingAutoencoder.n_samples = {}'.format(n_samples),\n",
        "    'oscillator_bank.use_angular_cumsum = True',  # Avoids cumsum accumulation errors.\n",
        "]\n",
        "\n",
        "with gin.unlock_config():\n",
        "  gin.parse_config(gin_params)\n",
        "\n",
        "\n",
        "\n",
        "# Set up the model just to predict audio given new conditioning\n",
        "model = ddsp.training.models.TranscribingAutoencoder()\n",
        "model.restore(ckpt)\n",
        "\n",
        "# Build model by running a batch through it.\n",
        "start_time = time.time()\n",
        "_ = model({'audio': audio}, training=False)\n",
        "print('Restoring model took %.1f seconds' % (time.time() - start_time))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "OBxFppRy9dDM"
      },
      "outputs": [],
      "source": [
        "#@title #Predict Pitch\n",
        "\n",
        "#@markdown Compare DDSP-INV (self-supervised) and [CREPE](https://github.com/marl/crepe) (supervised) models\n",
        "\n",
        "# DDSP-INV,\n",
        "start_time = time.time()\n",
        "print('\\nExtracting f0 with DDSP-INV...')\n",
        "controls = model.get_controls({'audio': audio}, training=False)\n",
        "print('Prediction took %.1f seconds' % (time.time() - start_time))\n",
        "\n",
        "# CREPE.\n",
        "start_time = time.time()\n",
        "print('\\nExtracting f0 with CREPE...')\n",
        "ddsp.spectral_ops.reset_crepe()\n",
        "f0_crepe, f0_confidence = ddsp.spectral_ops.compute_f0(audio[0], \n",
        "                                                       sample_rate=16000,\n",
        "                                                       frame_rate=31.25,\n",
        "                                                       viterbi=False)\n",
        "print('Prediction took %.1f seconds' % (time.time() - start_time))\n",
        "\n",
        "# Synthesize the CREPE audio\n",
        "synth = ddsp.synths.Wavetable(n_samples=n_samples, scale_fn=None)\n",
        "wavetable = np.sin(np.linspace(0, 2.0 * np.pi, 2048))[np.newaxis, np.newaxis, :]\n",
        "amps = np.ones([1, time_steps, 1]) * 0.1\n",
        "audio_crepe = synth(amps, wavetable, f0_crepe[np.newaxis, :, np.newaxis])\n",
        "\n",
        "# Synthesize the DDSP-INV audio\n",
        "audio_ddsp_inv = synth(controls['harm_amp'], wavetable, controls['f0_hz'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "jyseo-7W95AL"
      },
      "outputs": [],
      "source": [
        "#@title #Plot Resynthesis\n",
        "\n",
        "k = 0\n",
        "\n",
        "# Plot Pitch.\n",
        "plt.figure(figsize=(6, 4))\n",
        "f0_crepe_midi = ddsp.core.hz_to_midi(f0_crepe)\n",
        "f0_harm_midi = ddsp.core.hz_to_midi(controls['f0_hz'])\n",
        "plt.plot(np.ravel(f0_crepe_midi), label='crepe')\n",
        "plt.plot(np.ravel(f0_harm_midi[k]), label='ddsp-inv')\n",
        "plt.ylabel('Pitch (MIDI)')\n",
        "plt.xlabel('Time')\n",
        "plt.xticks([])\n",
        "plt.legend(loc='upper right')\n",
        "\n",
        "# Play Audio.\n",
        "print('Original')\n",
        "play(audio)\n",
        "\n",
        "print('Sinusoidal Resynthesis')\n",
        "play(controls['sin_audio'][k])\n",
        "\n",
        "print('Harmonic Resynthesis')\n",
        "play(controls['harm_audio'][k])\n",
        "\n",
        "print('DDSP-INV Pitch')\n",
        "play(audio_ddsp_inv[k])\n",
        "\n",
        "print('CREPE Pitch')\n",
        "play(audio_crepe[k])\n",
        "\n",
        "# Plot spectrograms.\n",
        "specplot(audio)\n",
        "plt.title(\"Original\")\n",
        "\n",
        "specplot(controls['sin_audio'][k])\n",
        "_ = plt.title(\"Sinusoidal Resynthesis\")\n",
        "\n",
        "specplot(controls['harm_audio'][k])\n",
        "_ = plt.title(\"Harmonic Resynthesis\")\n",
        "\n",
        "\n",
        "# Plot sinusoids.\n",
        "plt.figure(figsize=(6, 6))\n",
        "t = np.arange(controls['sin_freqs'].shape[1])\n",
        "for a, f in zip(np.transpose(controls['sin_amps'][k]), np.transpose(controls['sin_freqs'][k])):\n",
        "  plt.scatter(t, f, s=a*200, linewidths=1)\n",
        "  plt.ylim(0, 8000)\n",
        "plt.title('Sinusoids (Sinusoidal)')\n",
        "plt.ylabel('Frequency (Hz)')\n",
        "plt.xlabel('Time')\n",
        "plt.xticks([])\n",
        "\n",
        "plt.figure(figsize=(6, 6))\n",
        "t = np.arange(controls['harm_freqs'].shape[1])\n",
        "for a, f in zip(np.transpose(controls['harm_amps'][k]), np.transpose(controls['harm_freqs'][k])):\n",
        "  plt.scatter(t, f, s=a*200, linewidths=1)\n",
        "  plt.ylim(0, 8000)\n",
        "plt.title('Sinusoids (Harmonic)')\n",
        "plt.ylabel('Frequency (Hz)')\n",
        "plt.xlabel('Time')\n",
        "_ = plt.xticks([])\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "3YLyiTwPfVCT"
      ],
      "last_runtime": {},
      "name": "pitch_detection.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
